use std::collections::HashMap;
use crate::rpc_server::rpc_server::{AgentInfo, LabelSelector, GroupInfo};

#[derive(Debug, Clone)]
pub struct LabelManager {
    // agent_id -> AgentInfo
    agents: HashMap<String, AgentInfo>,
    // group_id -> GroupInfo  
    groups: HashMap<String, GroupInfo>,
    // label_key -> {label_value -> [agent_ids]}
    label_index: HashMap<String, HashMap<String, Vec<String>>>,
}

impl LabelManager {
    pub fn new() -> Self {
        Self {
            agents: HashMap::new(),
            groups: HashMap::new(),
            label_index: HashMap::new(),
        }
    }

    /// 注册Agent并建立标签索引
    pub fn register_agent(&mut self, agent_info: AgentInfo) -> Vec<String> {
        let agent_id = agent_info.agent_id.clone();
        
        // 清理旧的标签索引
        if let Some(old_agent) = self.agents.get(&agent_id).cloned() {
            self.remove_agent_from_label_index(&agent_id, &old_agent.labels);
        }
        
        // 建立新的标签索引
        for (key, value) in &agent_info.labels {
            self.label_index
                .entry(key.clone())
                .or_insert_with(HashMap::new)
                .entry(value.clone())
                .or_insert_with(Vec::new)
                .push(agent_id.clone());
        }
        
        // 自动分配到匹配的组
        let assigned_groups = self.auto_assign_to_groups(&agent_info);
        
        self.agents.insert(agent_id, agent_info);
        assigned_groups
    }

    /// 从标签索引中移除Agent
    fn remove_agent_from_label_index(&mut self, agent_id: &str, labels: &HashMap<String, String>) {
        for (key, value) in labels {
            if let Some(value_map) = self.label_index.get_mut(key) {
                if let Some(agent_list) = value_map.get_mut(value) {
                    agent_list.retain(|id| id != agent_id);
                    if agent_list.is_empty() {
                        value_map.remove(value);
                    }
                }
                if value_map.is_empty() {
                    self.label_index.remove(key);
                }
            }
        }
    }

    /// 根据标签选择器匹配Agent
    pub fn match_agents(&self, selector: &LabelSelector) -> Vec<String> {
        let mut matched_agents = Vec::new();
        
        for (agent_id, agent_info) in &self.agents {
            if self.agent_matches_selector(agent_info, selector) {
                matched_agents.push(agent_id.clone());
            }
        }
        
        matched_agents
    }

    /// 检查Agent是否匹配选择器
    fn agent_matches_selector(&self, agent: &AgentInfo, selector: &LabelSelector) -> bool {
        // 检查所有标签是否匹配
        for (key, value) in &selector.match_labels {
            if let Some(agent_value) = agent.labels.get(key) {
                if agent_value != value {
                    return false;
                }
            } else {
                return false;
            }
        }
        true
    }

    /// 自动分配Agent到匹配的组
    fn auto_assign_to_groups(&self, agent_info: &AgentInfo) -> Vec<String> {
        let mut assigned_groups = Vec::new();
        
        for (group_id, group_info) in &self.groups {
            if let Some(selector) = &group_info.selector {
                if self.agent_matches_selector(agent_info, selector) {
                    assigned_groups.push(group_id.clone());
                }
            }
        }
        
        assigned_groups
    }

    /// 创建标签组
    pub fn create_label_group(&mut self, group_name: String, selector: LabelSelector) -> String {
        let uuid_str = uuid::Uuid::new_v4().to_string();
        let group_id = format!("group_{}", &uuid_str[..8]);
        let matched_agents = self.match_agents(&selector);
        
        let group_info = GroupInfo {
            group_id: group_id.clone(),
            group_name,
            selector: Some(selector),
            agent_ids: matched_agents,
            created_at: chrono::Utc::now().timestamp(),
        };
        
        self.groups.insert(group_id.clone(), group_info);
        group_id
    }

    /// 获取Agent信息
    pub fn get_agent(&self, agent_id: &str) -> Option<&AgentInfo> {
        self.agents.get(agent_id)
    }

    /// 获取所有Agent
    pub fn get_all_agents(&self) -> Vec<&AgentInfo> {
        self.agents.values().collect()
    }

    /// 获取所有组
    pub fn get_all_groups(&self) -> Vec<&GroupInfo> {
        self.groups.values().collect()
    }

    /// 获取组信息
    pub fn get_group(&self, group_id: &str) -> Option<&GroupInfo> {
        self.groups.get(group_id)
    }

    /// 删除Agent
    pub fn remove_agent(&mut self, agent_id: &str) {
        if let Some(agent_info) = self.agents.remove(agent_id) {
            self.remove_agent_from_label_index(agent_id, &agent_info.labels);
            
            // 从所有组中移除该Agent
            for group_info in self.groups.values_mut() {
                group_info.agent_ids.retain(|id| id != agent_id);
            }
        }
    }

    /// 更新组的Agent列表
    pub fn update_group_agents(&mut self, group_id: &str) {
        if let Some(group_info) = self.groups.get(group_id) {
            if let Some(selector) = group_info.selector.clone() {
                let matched_agents = self.match_agents(&selector);
                if let Some(group_info) = self.groups.get_mut(group_id) {
                    group_info.agent_ids = matched_agents;
                }
            }
        }
    }

    /// 获取Agent的统计信息
    pub fn get_agent_stats(&self) -> HashMap<String, usize> {
        let mut stats = HashMap::new();
        
        // 按标签统计
        for (key, value_map) in &self.label_index {
            for (value, agent_list) in value_map {
                let label_key = format!("{}={}", key, value);
                stats.insert(label_key, agent_list.len());
            }
        }
        
        stats.insert("total_agents".to_string(), self.agents.len());
        stats.insert("total_groups".to_string(), self.groups.len());
        
        stats
    }
}